#pragma once

#include <algorithm>
#include <numeric>
#include <cstddef>

namespace mdtx::online::algo
{
    /**
     * @brief Weighted high median
     *
     */
    class whimed
    {
    public:
        /**
         * @brief Run the algo
         *
         * @details finds weighted high median using binary searh, input is
         * CHANGED IN-PLACE, so no allocation is made.
         *
         * @tparam iter input iterator type, point to pair/tuple of (value, weight)
         * @param it iterator to begin of input data
         * @param n number of data
         * @return auto
         */
        template <typename iter>
        static auto run(iter it, ptrdiff_t n)
        {
            using elem_type = typename std::iterator_traits<iter>::value_type;
            using value_type = std::tuple_element_t<0, elem_type>;
            using weight_type = std::tuple_element_t<1, elem_type>;

            auto sum_weight = [](const auto &sum, const auto &elem)
            {
                return sum + elem.second;
            };

            auto value = [it](ptrdiff_t idx) -> decltype(auto)
            {
                return std::get<0>(*(it + idx));
            };

            auto weight = [it](ptrdiff_t idx) -> decltype(auto)
            {
                return std::get<1>(*(it + idx));
            };

            value_type trial;
            weight_type w_half, w_left, w_right;
            w_half = std::accumulate(it, it + n, weight_type{}, sum_weight) / 2;

            ptrdiff_t left = 0, right = n - 1;
            while (true)
            {
                auto mid = right - left;
                if (mid == 0)
                {
                    return value(left);
                }
                else if (mid == 1)
                {
                    if (weight(left) >= weight(right))
                    {
                        return value(left);
                    }
                    else
                    {
                        return value(right);
                    }
                }

                // bsearch, mid is set right - left
                mid = left + mid / 2;
                std::nth_element(it + left, it + mid, it + right + 1,
                                 [](const auto &lhs, const auto &rhs)
                                 {
                                     return lhs.first < rhs.first;
                                 });
                trial = value(mid);

                w_left = std::accumulate(it + left, it + mid, weight_type{}, sum_weight);
                w_right = std::accumulate(it + mid + 1, it + right + 1, weight_type{}, sum_weight);

                if (w_left < w_half && w_right < w_half)
                {
                    return trial;
                }
                else if (w_left > w_half)
                {
                    weight(mid) += w_right;
                    right = mid;
                }
                else
                {
                    weight(mid) += w_left;
                    left = mid;
                }
            }
        }
    };
}
